module poly_square_mod
  use recon_type_mod
  use recon_math_mod
  use poly_utils_mod

  implicit none
  
contains

  subroutine poly_square_init(this, nd, sw, swx, swy, swz, xc, yc, zc, dx, dy, dz, ic, jc, kc, is, ie, js, je, ks, ke, mask, id)

    class(recon_type), intent(inout) :: this
    integer , intent(in), optional :: nd
    integer , intent(in), optional :: sw
    integer , intent(in), optional :: swx
    integer , intent(in), optional :: swy
    integer , intent(in), optional :: swz
    real(16), intent(in), optional :: xc(:)
    real(16), intent(in), optional :: yc(:)
    real(16), intent(in), optional :: zc(:)
    real(16), intent(in), optional :: dx
    real(16), intent(in), optional :: dy
    real(16), intent(in), optional :: dz
    integer , intent(in), optional :: ic
    integer , intent(in), optional :: jc
    integer , intent(in), optional :: kc
    integer , intent(in), optional :: is
    integer , intent(in), optional :: ie
    integer , intent(in), optional :: js
    integer , intent(in), optional :: je
    integer , intent(in), optional :: ks
    integer , intent(in), optional :: ke
    integer , intent(in), optional :: mask(:,:,:) ! Cell mask
    integer , intent(in), optional :: id

    integer i, j, k

    call this%clear()

    if(nd/=2)then
      stop 'poly_square_init support only nd==2'
    endif
    
    this%sw  = sw
    this%swx = sw
    this%swy = sw
    this%swz = 1
    this%nd  = nd
    this%nc  = sw**nd
    
    this%dx = dx
    this%dy = dy
    if(present(dz))then
      this%dz = dz
    else
      this%dz = 1
    endif

    if(present(is))then; this%is = is; else; this%is = 1       ; endif
    if(present(ie))then; this%ie = ie; else; this%ie = this%swx; endif
    if(present(js))then; this%js = js; else; this%js = 1       ; endif
    if(present(je))then; this%je = je; else; this%je = this%swy; endif
    if(present(ks))then; this%ks = ks; else; this%ks = 1       ; endif
    if(present(ke))then; this%ke = ke; else; this%ke = 1       ; endif

    allocate(this%xc(this%is:this%ie))
    allocate(this%yc(this%js:this%je))
    if (present(xc) .and. present(yc)) then
      this%xc = xc
      this%yc = yc
    else
      ! Set coordinates of cells on the large stencil with origin at center.
      do i = this%is, this%ie
        this%xc(i) = -int(sw / 2) + i - 1
      end do
      do j = this%js, this%je
        this%yc(j) = -int(sw / 2) + j - 1
      end do
    end if

    this%initialized = .true.

  end subroutine poly_square_init

  subroutine poly_square_calc_recon_matrix(this, ierr)

    class(recon_type), intent(inout) :: this
    integer, intent(out) :: ierr

    ! Local double double arrays for preserving precision.
    real(16), allocatable, dimension(:,:) :: A, iA
    integer i, j, k, ipt, n, iterm, dxn, dyn

    ierr = 0

    if (allocated(this%iA        )) deallocate(this%iA        )
    if (allocated(this%poly      )) deallocate(this%poly      )
    if (allocated(this%poly_r16  )) deallocate(this%poly_r16  )
    if (allocated(this%recon_mtx )) deallocate(this%recon_mtx )
    if (allocated(this%dpoly     )) deallocate(this%dpoly     )
    if (allocated(this%dpoly_r16 )) deallocate(this%dpoly_r16 )
    if (allocated(this%drecon_mtx)) deallocate(this%drecon_mtx)

    allocate(this%iA        (this%nc ,this%nc))
    allocate(this%poly      (this%npt,this%nc))
    allocate(this%poly_r16  (this%npt,this%nc))
    allocate(this%recon_mtx (this%npt,this%nc))
    allocate(this%dpoly     (this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1))
    allocate(this%dpoly_r16 (this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1))
    allocate(this%drecon_mtx(this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1))

    allocate( A(this%nc,this%nc))
    allocate(iA(this%nc,this%nc))

    ! Set the poly for each evaluation point.
    ! Select monomials according to mask.
    do ipt = 1, this%npt
      iterm = 1
      do j = 1, this%sw
        do i = 1, this%sw
          call calc_monomial(this%x(ipt), this%y(ipt), i - 1, j - 1, this%poly_r16(ipt,iterm))
          do dyn = 0, this%swy - 1
            do dxn = 0, this%swx - 1
              call calc_deriv_monomial(this%x(ipt), this%y(ipt), i - 1, j - 1, dxn, dyn, this%swx, this%swy, this%dpoly_r16(ipt,iterm,dxn,dyn,0))
              this%dpoly_r16(ipt,iterm,dxn,dyn,0) = this%dpoly_r16(ipt,iterm,dxn,dyn,0) / ( this%dx**dxn * this%dy**dyn )
            end do
          end do
          iterm = iterm + 1
        end do
      end do
    end do
    
    ! Calculate inverse of integral coefficient matrix.
    call calc_tensor_product_poly_integral_matrix(this%sw, this%sw, this%xc, this%yc, A)
    call inverse_matrix(A, iA, ierr)
    if (ierr /= 0) then
      deallocate(A, iA)
      return
    end if
    
    this%iA    = iA
    this%poly  = this%poly_r16
    this%dpoly = this%dpoly_r16
    
    this%recon_mtx  = matmul(this%poly_r16 , iA)
    
    do dyn = 0, this%swy - 1
      do dxn = 0, this%swx - 1
        this%drecon_mtx(:,:,dxn,dyn,0) = matmul(this%dpoly_r16(:,:,dxn,dyn,0), iA)
      end do
    end do
    
    deallocate(A, iA)

  end subroutine poly_square_calc_recon_matrix
  
  subroutine poly_square_reconstruct(this, fi, fo, TCI, ierr)

    class(recon_type), intent(inout) :: this
    real(8), intent(in ) :: fi(:,:,:) ! Cell averaged function values
    real(8), intent(out) :: fo(:)     ! Reconstructed function values on evaluation points
    logical, intent(in ), optional :: TCI ! Trouble Cell Indicator(TCI), 1 for existing TC(use WENO), 0 for no TC(use poly)
    integer, intent(out), optional :: ierr

#ifdef DEBUG
    if(present(ierr))then
      ierr = 0
      if (size(fi) /= this%nc) then
        if (present(ierr)) ierr = 1
        return
      end if
      if (size(fo) /= this%npt) then
        if (present(ierr)) ierr = 1
        return
      end if
    endif
#endif

    fo = matmul(this%recon_mtx, pack(fi, .true.))
    
  end subroutine poly_square_reconstruct
  
  subroutine poly_square_reconstruct_deriv(this, fi, fx, fy, fz, ierr)

    class(recon_type), intent(inout) :: this
    real(8), intent(in ) :: fi(:,:,:) ! Cell averaged function values
    real(8), intent(out), optional :: fx(:)     ! Reconstructed x derivative values on evaluation points
    real(8), intent(out), optional :: fy(:)     ! Reconstructed y derivative values on evaluation points
    real(8), intent(out), optional :: fz(:)     ! Reconstructed y derivative values on evaluation points
    integer, intent(out), optional :: ierr

#ifdef DEBUG
    if(present(ierr))then
      ierr = 0
      if (size(fi) /= this%nc) then
        if (present(ierr)) ierr = 1
        return
      end if
      if (size(fo) /= this%npt) then
        if (present(ierr)) ierr = 1
        return
      end if
    endif
#endif

    fx = matmul(this%drecon_mtx(:,:,1,0,0), pack(fi, .true.))
    fy = matmul(this%drecon_mtx(:,:,0,1,0), pack(fi, .true.))
    
  end subroutine poly_square_reconstruct_deriv

end module poly_square_mod
